#!/bin/bash

# 强制错误检查和未定义变量检测
set -euo pipefail

# 配置参数
WG_NET="10.252.0.0/24"
SERVER_IP="10.252.0.1"
SERVER_PORT=8051
SERVER_PUBLIC_ADDR="1.180.177.139"  # 客户端连接地址
FIREWALL_IP="192.168.100.87"             # 防火墙公网IP
DNS_SERVER="8.8.8.8"
CONFIG_FILE="/etc/wireguard/wg0.conf"
CLIENT_DIR="/etc/wireguard/clients"
NEXT_IP_FILE="/etc/wireguard/next_ip"

# 检查root权限
if [ "$EUID" -ne 0 ]; then
    echo "请使用sudo或root用户运行此脚本"
    exit 1
fi

# 清理旧配置
clean_legacy() {
    echo "[清理] 移除旧配置..."
    systemctl stop wg-quick@wg0 2>/dev/null || true
    systemctl disable wg-quick@wg0 2>/dev/null || true
    rm -f /etc/wireguard/{wg0.conf,server_privatekey,server_publickey} 2>/dev/null
    rm -rf "$CLIENT_DIR" 2>/dev/null
    ufw delete allow ${SERVER_PORT}/udp 2>/dev/null || true
}

# 域名解析验证
validate_dns() {
    if [[ $SERVER_PUBLIC_ADDR =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]; then
        echo "[信息] 使用IP地址：$SERVER_PUBLIC_ADDR"
        return 0
    fi

    echo "[信息] 正在解析域名：$SERVER_PUBLIC_ADDR"
    
    local resolved_ip=""
    if command -v dig &>/dev/null; then
        resolved_ip=$(dig +short "$SERVER_PUBLIC_ADDR" | head -n1)
    elif command -v nslookup &>/dev/null; then
        resolved_ip=$(nslookup "$SERVER_PUBLIC_ADDR" | awk '/^Address: / { print $2 }' | tail -n1)
    fi

    if [[ ! "$resolved_ip" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]; then
        echo "[错误] 域名解析失败: $SERVER_PUBLIC_ADDR"
        return 1
    fi
    
    echo "[成功] 域名解析到IP：$resolved_ip (客户端仍使用域名连接)"
}

# 生成客户端配置
generate_client() {
    [ -f "$NEXT_IP_FILE" ] && next_ip=$(cat "$NEXT_IP_FILE") || next_ip=2
    CLIENT_IP="10.252.0.$next_ip"
    CLIENT_NAME="client_$next_ip"
    CLIENT_FILE="$CLIENT_DIR/$CLIENT_NAME.conf"
    
    # 生成密钥对
    CLIENT_PRIVATEKEY=$(wg genkey)
    CLIENT_PUBLICKEY=$(echo "$CLIENT_PRIVATEKEY" | wg pubkey)

    # 更新服务端配置
    wg set wg0 peer "$CLIENT_PUBLICKEY" allowed-ips "$CLIENT_IP/32"
    wg-quick save wg0 >/dev/null 2>&1
    
    # 生成客户端文件（强制使用域名）
    cat > "$CLIENT_FILE" <<-EOF
[Interface]
PrivateKey = $CLIENT_PRIVATEKEY
Address = $CLIENT_IP/24
DNS = $DNS_SERVER

[Peer]
PublicKey = $(cat /etc/wireguard/server_publickey)
Endpoint = ${SERVER_PUBLIC_ADDR}:80
AllowedIPs = 0.0.0.0/0
PersistentKeepalive = 25
EOF

    # 生成二维码
    qrencode -t png -o "$CLIENT_DIR/$CLIENT_NAME.png" -r "$CLIENT_FILE"
    ((next_ip++))
    [ $next_ip -lt 254 ] && echo $next_ip > "$NEXT_IP_FILE" || rm -f "$NEXT_IP_FILE"
    
    echo "客户端配置已保存至：$CLIENT_FILE"
    echo "二维码文件：$CLIENT_DIR/$CLIENT_NAME.png"
}

# 主安装流程
install_wireguard() {
    clean_legacy
    
    # 安装依赖
    echo "[系统] 更新软件源..."
    apt-get update -qq

    echo "[安装] 安装必要组件..."
    DEBIAN_FRONTEND=noninteractive apt-get install -y -qq \
        wireguard \
        wireguard-tools \
        qrencode \
        ufw \
        dnsutils

    # 验证域名解析
    if ! validate_dns; then
        exit 1
    fi

    # 生成密钥对
    echo "[配置] 生成密钥..."
    umask 077
    mkdir -p /etc/wireguard
    wg genkey | tee /etc/wireguard/server_privatekey | wg pubkey > /etc/wireguard/server_publickey

    # 创建服务端配置
    DEFAULT_IFACE=$(ip route | awk '/default/ {print $5; exit}')
    cat > "$CONFIG_FILE" <<-EOF
[Interface]
Address = $SERVER_IP/24
ListenPort = $SERVER_PORT
PrivateKey = $(cat /etc/wireguard/server_privatekey)
PostUp = iptables -A FORWARD -i wg0 -j ACCEPT; iptables -t nat -A POSTROUTING -o $DEFAULT_IFACE -j MASQUERADE
PostDown = iptables -D FORWARD -i wg0 -j ACCEPT; iptables -t nat -D POSTROUTING -o $DEFAULT_IFACE -j MASQUERADE
EOF

    # 防火墙配置
    echo "[防火墙] 配置规则..."
    ufw --force enable || true
    ufw allow from 0.0.0.0/0 to any port $SERVER_PORT proto udp
    ufw allow ${SERVER_PORT}/udp
    ufw --force reload

    # 启用IP转发
    echo "[系统] 启用IP转发..."
    sed -i '/net.ipv4.ip_forward/d' /etc/sysctl.conf
    echo "net.ipv4.ip_forward=1" >> /etc/sysctl.conf
    sysctl -p >/dev/null

    # 启动服务
    echo "[服务] 启动WireGuard..."
    systemctl enable --now wg-quick@wg0

    # 初始化客户端
    mkdir -p "$CLIENT_DIR"
    echo 2 > "$NEXT_IP_FILE"
    generate_client
}

# --------------------------
# 卸载WireGuard主服务
# --------------------------
uninstall_wg() {
    echo -e "\033[31m[卸载] 正在卸载WireGuard...\033[0m"
    
    # 停止并禁用服务
    systemctl stop wg-quick@wg0 2>/dev/null || true
    systemctl disable wg-quick@wg0 2>/dev/null || true

    # 删除配置文件和密钥
    rm -f /etc/wireguard/{wg0.conf,server_privatekey,server_publickey} 2>/dev/null
    rm -rf "$CLIENT_DIR" 2>/dev/null
    
    # 清理防火墙规则
    ufw delete allow $LOCAL_LISTEN_PORT/udp 2>/dev/null || true
    ufw --force reload
    
    # 禁用IP转发（如果这是唯一使用它的服务）
    sed -i '/net.ipv4.ip_forward/d' /etc/sysctl.conf
    sysctl -p >/dev/null

    # 移除WireGuard相关软件包
    DEBIAN_FRONTEND=noninteractive apt-get purge -y wireguard wireguard-tools qrencode 2>/dev/null
    
    echo -e "\033[31m[完成] WireGuard已卸载！\033[0m"
}

# --------------------------
# 执行逻辑
# --------------------------
if [[ "$*" == "--force" ]]; then
    echo "[模式] 强制重新安装..."
    install_wireguard
elif systemctl is-active wg-quick@wg0 >/dev/null 2>&1; then
    echo "[模式] 生成新客户端配置..."
    validate_dns
    generate_client
elif [ -f "$CONFIG_FILE" ]; then
    echo "[模式] 检测到残留配置，重新初始化..."
    install_wireguard
else
    echo "[模式] 全新安装..."
    install_wireguard
fi


echo -e "\033[32m[完成] 部署成功！验证命令：wg show wg0\033[0m"